package com.syf.starter.core.netty.handle;

import cn.hutool.core.util.StrUtil;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelId;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.ReferenceCountUtil;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
@ChannelHandler.Sharable
public class PreventDoSAttackHandle extends ChannelInboundHandlerAdapter {
    private final Map<ChannelId, RequestCounter> REQUEST_COUNTS_MAP = new ConcurrentHashMap<>();

    private final Long TIME_WINDOW;

    private final Integer DOS_ATTACK_THRESHOLD;

    public PreventDoSAttackHandle(Integer DOS_ATTACK_THRESHOLD, String DOS_CALS_TIME_WINDOW_UNIT) {
        super();
        if (StrUtil.equals(DOS_CALS_TIME_WINDOW_UNIT, "2")) {
            TIME_WINDOW = TimeUnit.MINUTES.toMillis(1);
        } else {
            TIME_WINDOW = TimeUnit.SECONDS.toMillis(1);
        }

        this.DOS_ATTACK_THRESHOLD = DOS_ATTACK_THRESHOLD;
    }


    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (isDoSAttack(ctx.channel().id())) {
            log.error("疑似DoS，客户端：[{}]，请求数据:{}", ctx.channel().remoteAddress().toString(), msg);
            ReferenceCountUtil.release(msg);
            ctx.channel().close();
            return;
        }

        ctx.fireChannelRead(msg);

    }

    private boolean isDoSAttack(ChannelId channelId) {
        RequestCounter requestCounter = REQUEST_COUNTS_MAP.computeIfAbsent(channelId, k -> new RequestCounter(TIME_WINDOW));
        if (requestCounter.incrementAndGet() > DOS_ATTACK_THRESHOLD) {
            return true;
        }

        return false;
    }

    public static class RequestCounter {
        private long timeStamp = System.currentTimeMillis();

        private Long TIME_WINDOW;

        private int count = 0;

        public RequestCounter(Long TIME_WINDOW) {
            this.TIME_WINDOW = TIME_WINDOW;
        }

        public synchronized int incrementAndGet() {
            long now = System.currentTimeMillis();
            if (now > timeStamp + TIME_WINDOW) {
                timeStamp = now;
                count = 0;
            }

            return count + 1;
        }
    }

}
